import json
from Cryptodome.Cipher import PKCS1_OAEP
from Cryptodome.PublicKey import RSA
from Cryptodome.Signature import pkcs1_15
from Cryptodome.Hash import SM3
import base64


class BankGateway:
    def __init__(self, private_key, public_key):
        # 加载私钥和公钥
        self.private_key = RSA.import_key(base64.b64decode(private_key))
        self.public_key = RSA.import_key(base64.b64decode(public_key))

    def sign_data(self, data):
        # 将数据转换为字节流
        data_bytes = data.encode('utf-8')

        # 计算SM3哈希值
        hash_obj = SM3.new(data_bytes)
        hash_value = hash_obj.digest()

        # 使用私钥进行签名
        signer = pkcs1_15.new(self.private_key)
        signature = signer.sign(hash_obj)

        # 返回base64编码后的签名结果
        return base64.b64encode(signature).decode('utf-8')

    def verify_signature(self, data, signature):
        # 将数据转换为字节流
        data_bytes = data.encode('utf-8')

        # 计算SM3哈希值
        hash_obj = SM3.new(data_bytes)
        hash_value = hash_obj.digest()

        # 验证签名
        verifier = pkcs1_15.new(self.public_key)
        try:
            verifier.verify(hash_obj, base64.b64decode(signature))
            return True
        except ValueError:
            return False

    def encrypt_data(self, data):
        # 将数据转换为字节流
        data_bytes = data.encode('utf-8')

        # 使用公钥进行加密
        cipher = PKCS1_OAEP.new(self.public_key)
        encrypted_data = cipher.encrypt(data_bytes)

        # 返回base64编码后的加密结果
        return base64.b64encode(encrypted_data).decode('utf-8')

    def decrypt_data(self, encrypted_data):
        # 解码base64编码的加密数据
        encrypted_bytes = base64.b64decode(encrypted_data)

        # 使用私钥进行解密
        cipher = PKCS1_OAEP.new(self.private_key)
        decrypted_data = cipher.decrypt(encrypted_bytes)

        # 返回解密后的数据
        return decrypted_data.decode('utf-8')

    def process_request(self, request):
        # 解析HTTP请求头
        headers = request.headers
        appid = headers.get('appid')
        timestamp = headers.get('timestamp')
        sign = headers.get('sign')
        verify = headers.get('verify')
        apisign = headers.get('apisign')

        # 构建待签名字符串
        sign_string = f"appid={appid}&secret={sign}&sign={sign}&timestamp={timestamp}"

        # 验证签名
        if verify == 'SM3withSM2' and self.verify_signature(sign_string, apisign):
            # 解析HTTP请求体
            body = json.loads(request.body)

            # 解密业务报文
            encrypted_data = body['data']
            decrypted_data = self.decrypt_data(encrypted_data)

            # 在此处可以对解密后的报文进行处理和业务逻辑操作
            # ...

            # 构建响应报文
            response_data = {
                'returnCode': 'SUC0000',
                'errorMsg': None,
                'data': None  # 可以填充需要透传的返回数据
            }
            response_body = json.dumps(response_data)
            return response_body
        else:
            # 签名验证失败
            response_data = {
                'returnCode': 'ERR0001',
                'errorMsg': 'Signature verification failed',
                'data': None
            }
            responseBody = json.dumps(response_data)
            return responseBody
